/* 
 * File:   MRFSegmentation.h
 * Author: hadrien
 *
 * Created on 18 avril 2012, 15:15
 */

#ifndef MRFSEGMENTATIONDIAG_H
#define	MRFSEGMENTATIONDIAG_H
#include <float.h>
#include <cv.h>
#include <cstdlib>
#include "SaveLabelImage.h"

template<typename OPTIMIZATION, typename EMPARAMETERS>
class MRFSegmentationDiag {
public:
    typedef OPTIMIZATION optimization_t;
    typedef typename optimization_t::Params Params;
    typedef typename Params::features_vector_t features_vector_t;
    typedef typename Params::label_t label_t;
    typedef EMPARAMETERS EMParams;
    typedef GaussianProcess<Params> model_t;

    MRFSegmentationDiag() {
    }
    MRFSegmentationDiag(const MRFSegmentationDiag& orig);

    virtual ~MRFSegmentationDiag() {
    }
private:

    void randomInitialization(label_t** labels, model_t &model) const {
        for (int i = 0; i < Params::getHeight(); ++i) {
            for (int j = 0; j < Params::getWidth(); ++j) {
                label_t lab = rand() % (Params::getL());
                labels[i][j] = lab;
                for (int l = 0; l < Params::getL(); ++l) {
                    model.setAPosterioriLabel(i, j, l, 0);
                }
                model.setAPosterioriLabel(i, j, lab, 1);
            }
        }
    }

    void kmeansInitialization(features_vector_t const * const * obs, label_t** labels, model_t &model) const {
        cv::Mat_<float> samples(Params::getWidth() * Params::getHeight(), Params::getN());
        for (int i = 0; i < Params::getHeight(); ++i) {
            for (int j = 0; j < Params::getWidth(); ++j) {
                for (int n = 0; n < Params::getN(); ++n) {
                    samples(i * Params::getWidth() + j) = obs[i][j](n);
                }
            }
        }
        cv::Mat_<int> label_kmeans(Params::getWidth() * Params::getHeight(), 1);

        cv::TermCriteria crit(cv::TermCriteria::MAX_ITER + cv::TermCriteria::EPS, 500, 0.001);
        features_vector_t centers(Params::getL(), Params::getN());

        cv::kmeans(samples, Params::getL(), label_kmeans, crit, 3, cv::KMEANS_PP_CENTERS, centers);

        for (int i = 0; i < Params::getHeight(); ++i) {
            for (int j = 0; j < Params::getWidth(); ++j) {
                label_t lab = label_kmeans(i * Params::getWidth() + j);
                labels[i][j] = lab;
                for (int l = 0; l < Params::getL(); ++l) {
                    model.setAPosterioriLabel(i, j, l, 0);
                }
                model.setAPosterioriLabel(i, j, lab, 1);
            }
        }
    }

    //Return the likelyhood

    double APosterioriEstimation(label_t const * const * labels, features_vector_t const * const * obs, model_t &model) const {
        double result = 0;
        double det[Params::getL()];
        features_vector_t inv_var[Params::getL()];
        for (int l = 0; l < Params::getL(); ++l) {
            inv_var[l] = features_vector_t::zeros(Params::getN(),1);
            det[l] = 1;
            for (int n = 0; n < Params::getN(); ++n) {
                det[n] *= model.getVariance(l)(n);
                inv_var[l](n) = 1 / model.getVariance(l)(n);
            }
        }
        for (int i = 0; i < Params::getHeight(); ++i) {
            for (int j = 0; j < Params::getWidth(); ++j) {
                double sum = 0;
                for (int l = 0; l < Params::getL(); ++l) {
                    double p = 0;
                    for(int n = 0; n < Params::getN(); ++n) {
                        double d = obs[i][j](n) - model.getMean(l)(n);
                        p += d*d*inv_var[l](n);
                    }
                    p = exp(-p);
                    p = p * 1 / sqrt(pow(2.0 * 3.141592653589793, Params::getN()) * det[l]) * model.getAPrioriLabel(l);
                    sum += p;
                    model.setAPosterioriLabel(i, j, l, p);
                }
                for(int l = 0; l < Params::getL(); ++l) {
                    model.setAPosterioriLabel(i, j, l, model.getAPosterioriLabel(i, j, l) / sum);
                }
                result += log(sum);
            }
        }

        return result / Params::getHeight() / Params::getWidth();
    }

    void maximization(features_vector_t const * const * obs, model_t &model) const {
        for (int l = 0; l < Params::getL(); ++l) {
            features_vector_t mean(Params::getN(), 1);
            features_vector_t variance(Params::getN(), 1);
            mean = features_vector_t::zeros(Params::getN(), 1);
            variance = features_vector_t::zeros(Params::getN(), 1);

            //Mean
            double K = 0;
            for (int i = 0; i < Params::getHeight(); ++i) {
                for (int j = 0; j < Params::getWidth(); ++j) {
                    mean += obs[i][j] * model.getAPosterioriLabel(i, j, l);
                    K += model.getAPosterioriLabel(i, j, l);
                }
            }
            mean = mean / K;
            //Variance
            for (int i = 0; i < Params::getHeight(); ++i) {
                for (int j = 0; j < Params::getWidth(); ++j) {
                    for (int n = 0; n < Params::getN(); ++n) {
                        double d = (obs[i][j](n) - mean(n));
                        variance(n) += d*d * model.getAPosterioriLabel(i, j, l);
                    }
                }
            }
            variance = variance / K;
            model.setMean(l, mean);
            model.setVariance(l, variance);
            model.setAPrioriLabel(l, K / Params::getHeight() / Params::getWidth());
        }
        for (int a = 0; a < Params::getL(); ++a) {
            std::cout << "Mean " << a << " :\n" << model.getMean(a) << std::endl;
            std::cout << "Variance " << a << " :\n" << model.getVariance(a) << std::endl;
        }
    }

public:

    void operator()(features_vector_t const * const* obs, label_t** labels) const {
        model_t model;
        optimization_t opt;
        kmeansInitialization(obs, labels, model);
        maximization(obs, model);
        //Start with M-STEP
        double oldlogLikelihood = -DBL_MAX;
        int it = 0;
        double difflogLikelihood;
        char buffer[50];
        //M-STEP performed

        sprintf(buffer, "init.jpg", it);
        SaveLabelImage::saveImage(labels, Params::getHeight(), Params::getWidth(), Params::getL(), buffer);

        do {
            //E-STEP
            opt(model, obs, labels);
            double newlogLikelihood = APosterioriEstimation(labels, obs, model);
            sprintf(buffer, "segmented it=%d.jpg", it);
            SaveLabelImage::saveImage(labels, Params::getHeight(), Params::getWidth(), Params::getL(), buffer);

            //M-STEP
            maximization(obs, model);

            difflogLikelihood = fabs(newlogLikelihood - oldlogLikelihood);
            oldlogLikelihood = newlogLikelihood;
            it++;

            std::cout << "--------------------\nDiff LogLikehood : " << difflogLikelihood << "\n---------------------" << std::endl;
        } while (difflogLikelihood > EMParams::getStopCriterion() && it < EMParams::getMaxIterations());
    }
};

#endif	/* MRFSEGMENTATION_H */

